﻿////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
// <copyright>Copyright 2008-2011 Andreas Huber Doenni</copyright>
// Distributed under the GNU General Public License version 2 (GPLv2).
// See accompanying file License.txt or copy at http://phuse.codeplex.com/license.
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

namespace Phuse.Net.Sockets
{
    using System;
    using System.Diagnostics.CodeAnalysis;
    using System.Net;
    using System.Net.Sockets;
    using System.Threading;

    using Phuse.Net.Nntp;
    
    /// <summary>Minimalistic implementation of System.Net.Sockets.TcpClient.</summary>
    /// <threadsafety static="true" instance="false"/>
    public sealed class TcpClient : IDisposable
    {
        private static int connectTimeout = 30000;
        private readonly CancellationToken cancellationToken;
        private SocketAsyncEventArgs sendReceiveArgs;
        private readonly AutoResetEvent completed;
        private Socket socket;
        private int sendTimeout = Timeout.Infinite;
        private int receiveTimeout = Timeout.Infinite;

        ////////////////////////////////////////////////////////////////////////////////////////////////////////////////

        /// <summary>Initializes a new instance of the TcpClient class and connects to <paramref name="port"/> on the
        /// host named <paramref name="hostName"/>.</summary>
        /// <exception cref="ArgumentException"><paramref name="hostName"/> equals <c>""</c>.</exception>
        /// <exception cref="ArgumentNullException"><paramref name="hostName"/> equals <c>null</c>.</exception>
        /// <exception cref="ArgumentOutOfRangeException"><paramref name="port"/> is not between
        /// <see cref="IPEndPoint.MinPort"/> and <see cref="IPEndPoint.MaxPort"/>.</exception>
        /// <exception cref="SocketException">An error occurred while connecting.</exception>
        public TcpClient(string hostName, int port, CancellationToken cancellationToken)
        {
            this.cancellationToken = cancellationToken;

            using (var argsGuard = ScopeGuard.Create(this.sendReceiveArgs = new SocketAsyncEventArgs()))
            using (var completedGuard = ScopeGuard.Create(this.completed = new AutoResetEvent(false)))
            {
                this.socket = Connect(hostName, port, this.completed, this.cancellationToken);
                completedGuard.Dismiss();
                argsGuard.Dismiss();
            }
        }

        /// <summary>Shuts the connection down.</summary>
        [SuppressMessage("Microsoft.Design", "CA1031:DoNotCatchGeneralExceptionTypes", Justification = "Dispose must never throw.")]
        public void Dispose()
        {
            if (this.socket != null)
            {
                try
                {
                    this.socket.Shutdown(SocketShutdown.Both);
                    this.socket.Dispose();
                    this.completed.Dispose();
                    this.sendReceiveArgs.Dispose();
                }
                catch
                {
                }
                finally
                {
                    this.socket = null;
                }
            }
        }

        /// <summary>Gets the size of the send buffer.</summary>
        /// <exception cref="ObjectDisposedException"><see cref="Dispose"/> has been called.</exception>
        public int SendBufferSize
        {
            get
            {
                this.AssertNotDisposed();
                return this.socket.SendBufferSize;
            }
        }

        /// <summary>Gets the size of the receive buffer.</summary>
        /// <exception cref="ObjectDisposedException"><see cref="Dispose"/> has been called.</exception>
        public int ReceiveBufferSize
        {
            get
            {
                this.AssertNotDisposed();
                return this.socket.ReceiveBufferSize;
            }
        }

        /// <summary>Gets or sets the amount of time, in milliseconds, the constructor will wait for a successful
        /// connection.</summary>
        /// <value>Default is 30000.</value>
        public static int ConnectTimeout
        {
            get { return Interlocked.Add(ref connectTimeout, 0); }
            set { Interlocked.Exchange(ref connectTimeout, value); }
        }

        /// <summary>Gets or sets the amount of time, in milliseconds, a <see cref="TcpClient"/> object will wait for a
        /// <see cref="Send"/> operation to complete successfully.</summary>
        /// <value>Default is <see cref="Timeout.Infinite"/>.</value>
        public int SendTimeout
        {
            get { return this.sendTimeout; }
            set { this.sendTimeout = value; }
        }

        /// <summary>Gets or sets the amount of time, in milliseconds, a <see cref="TcpClient"/> object will wait for a
        /// <see cref="Receive"/> operation to complete successfully.</summary>
        /// <value>Default is <see cref="Timeout.Infinite"/>.</value>
        public int ReceiveTimeout
        {
            get { return this.receiveTimeout; }
            set { this.receiveTimeout = value; }
        }

        /// <summary>Gets a value that indicates whether the <see cref="TcpClient"/> is connected to a remote host as of
        /// the last <see cref="Send"/> or <see cref="Receive"/> operation.</summary>
        /// <exception cref="ObjectDisposedException"><see cref="Dispose"/> has been called.</exception>
        public bool Connected
        {
            get
            {
                this.AssertNotDisposed();
                return this.socket.Connected;
            }
        }

        /// <summary>Sends <paramref name="count"/> bytes from <paramref name="buffer"/>, starting at
        /// <paramref name="offset"/>.</summary>
        /// <exception cref="ArgumentException">The sum of <paramref name="offset"/> and <paramref name="count"/> is
        /// greater than the length of <paramref name="buffer"/>.</exception>
        /// <exception cref="ArgumentNullException"><paramref name="buffer"/> equals <c>null</c>.</exception>
        /// <exception cref="ArgumentOutOfRangeException"><paramref name="offset"/> and/or <paramref name="count"/> are
        /// negative.</exception>
        /// <exception cref="ObjectDisposedException"><see cref="Dispose"/> has been called.</exception>
        /// <exception cref="SocketException">An error occurred when attempting to access the socket.</exception>
        public void Send(byte[] buffer, int offset, int count)
        {
            this.AssertNotDisposed();

            // This optimization is essential for correct function on Windows Phone, as the socket there refuses to
            // accept a buffer with count == 0.
            if (count > 0)
            {
                var actualCount = Transfer(
                    ref this.sendReceiveArgs,
                    buffer,
                    offset,
                    count,
                    this.completed,
                    this.sendTimeout,
                    this.cancellationToken,
                    this.socket.SendAsync);

                if (actualCount != count)
                {
                    throw new SocketException();
                }
            }
        }

        /// <summary>Receives up to <paramref name="count"/> bytes into <paramref name="buffer"/>, starting at
        /// <paramref name="offset"/>.</summary>
        /// <returns>The number of bytes received.</returns>
        /// <exception cref="ArgumentException">The sum of <paramref name="offset"/> and <paramref name="count"/> is
        /// greater than the length of <paramref name="buffer"/>.</exception>
        /// <exception cref="ArgumentNullException"><paramref name="buffer"/> equals <c>null</c>.</exception>
        /// <exception cref="ArgumentOutOfRangeException"><paramref name="offset"/> and/or <paramref name="count"/> are
        /// negative.</exception>
        /// <exception cref="ObjectDisposedException"><see cref="Dispose"/> has been called.</exception>
        /// <exception cref="SocketException">An error occurred when attempting to access the socket.</exception>
        public int Receive(byte[] buffer, int offset, int count)
        {
            this.AssertNotDisposed();
            return Transfer(
                ref this.sendReceiveArgs,
                buffer,
                offset,
                count,
                this.completed,
                this.receiveTimeout,
                this.cancellationToken,
                this.socket.ReceiveAsync);
        }

        ////////////////////////////////////////////////////////////////////////////////////////////////////////////////

        private void AssertNotDisposed()
        {
            if (this.socket == null)
            {
                throw new ObjectDisposedException(this.GetType().ToString());
            }
        }

        private static Socket Connect(
            string hostName, int port, AutoResetEvent completed, CancellationToken cancellationToken)
        {
            object padlock = new object();

            try
            {
                if (hostName == null)
                {
                    throw new ArgumentNullException("hostName");
                }

                if ((port < IPEndPoint.MinPort) || (port > IPEndPoint.MaxPort))
                {
                    throw new ArgumentOutOfRangeException("port");
                }

                using (var args = new SocketAsyncEventArgs())
                {
                    args.RemoteEndPoint = new DnsEndPoint(hostName, port);
                    args.Completed += (sender, e) => OnCompleted(padlock, ref completed);

                    if (Socket.ConnectAsync(SocketType.Stream, ProtocolType.Tcp, args))
                    {
                        WaitAndThrowIfCanceledOrTimedOut(completed, ConnectTimeout, cancellationToken);
                    }

                    if (args.ConnectByNameError != null)
                    {
                        throw args.ConnectByNameError;
                    }

                    if (args.SocketError != SocketError.Success)
                    {
                        throw new SocketException((int)args.SocketError);
                    }

                    return args.ConnectSocket;
                }
            }
            finally
            {
                lock (padlock)
                {
                    completed = null;
                }
            }
        }

        private static int Transfer(
            ref SocketAsyncEventArgs args,
            byte[] buffer,
            int offset,
            int count,
            AutoResetEvent completed,
            int timeout,
            CancellationToken cancellationToken,
            Func<SocketAsyncEventArgs, bool> transferCallback)
        {
            BufferHelper.AssertValidRange(buffer, "buffer", offset, "offset", count, "count");
            object padlock = new object();

            try
            {
                args.SetBuffer(buffer, offset, count);
                EventHandler<SocketAsyncEventArgs> completedHandler = (s, e) => OnCompleted(padlock, ref completed);
                args.Completed += completedHandler;

                if (transferCallback(args))
                {
                    WaitAndThrowIfCanceledOrTimedOut(completed, timeout, cancellationToken);
                }

                if (args.SocketError != SocketError.Success)
                {
                    throw new SocketException((int)args.SocketError);
                }

                // Removing the handler is not necessary when an exception is thrown, because the SocketAsyncEventArgs
                // object will not be reused then, see catch block.
                args.Completed -= completedHandler;
                return args.BytesTransferred;
            }
            catch
            {
                // Some of what could possibly go wrong in the try block will very likely leave the SocketAsyncEventArgs
                // object in an unusable state. For example, when a timeout happens, the Completed event may be raised
                // *after* we have left this method. In this case it is dangerous to reuse the args object, because the
                // event might be raised when we have reentered this method. Note that unsubscribing the handler here
                // does not help at all, since the issue arises when this method is reentered and a new handler is added
                // for the Completed event.
                // Although the SocketAsyncEventArgs object is likely reusable for a certain class of errors, it does
                // not seem to be worth the additional testing effort to make sure. Note that most (if not all)
                // exceptions originating in this method will lead to the disposal of the TcpClient object anyway, so
                // performance will not even improve from the additional reuse of the SocketAsyncEventArgs object.
                args.Dispose();
                args = new SocketAsyncEventArgs();
                throw;
            }
            finally
            {
                lock (padlock)
                {
                    completed = null;
                }
            }
        }

        private static void OnCompleted(object padlock, ref AutoResetEvent completed)
        {
            lock (padlock)
            {
                if (completed != null)
                {
                    completed.Set();
                }
            }
        }

        private static void WaitAndThrowIfCanceledOrTimedOut(
            AutoResetEvent completed, int timeout, CancellationToken cancellationToken)
        {
            var waitResult = WaitHandle.WaitAny(new[] { completed, cancellationToken.WaitHandle }, timeout);

            if (waitResult == 1)
            {
                cancellationToken.ThrowIfCancellationRequested();
            }
            else if (waitResult == WaitHandle.WaitTimeout)
            {
                throw new SocketException((int)SocketError.TimedOut);
            }
        }
    }
}
